"""
PETSc-Only Runner Implementation.

This runner executes pure PETSc solvers using mpiexec.
"""

import os
import subprocess
import time
import logging
from pathlib import Path
from typing import Dict, List, Optional, Any

from gcr_solver_manager.extensions.base_runner import (
    RunnerInterface, RunConfig, RunResult, EnvironmentResult
)

logger = logging.getLogger(__name__)


class PETScOnlyRunner(RunnerInterface):
    """Runner for PETSc-Only framework."""
    
    def __init__(self, source_dir: Path):
        """Initialize runner with source directory."""
        self.source_dir = source_dir
    
    def run_solver(self, run_config: RunConfig) -> RunResult:
        """Execute solver with given configuration."""
        start_time = time.time()
        
        try:
            # Build mpiexec command for PETSc
            cmd = self._build_mpiexec_command(run_config)
            
            logger.info(f"Executing PETSc solver: {' '.join(cmd)}")
            
            # Set up environment with PETSc variables
            env = os.environ.copy()
            env.update(run_config.environment)
            
            # Ensure PETSc environment is properly set
            if 'PETSC_DIR' not in env:
                logger.warning("PETSC_DIR not set in environment")
            
            # Execute solver
            result = subprocess.run(
                cmd,
                cwd=self.source_dir,
                capture_output=True,
                text=True,
                env=env,
                timeout=run_config.timeout or 1800  # Default 30 minutes
            )
            
            duration = time.time() - start_time
            success = result.returncode == 0
            
            # Extract PETSc-specific metrics from output
            metrics = self._extract_petsc_metrics(result.stdout, result.stderr)
            
            run_result = RunResult(
                success=success,
                exit_code=result.returncode,
                stdout=result.stdout,
                stderr=result.stderr,
                duration=duration,
                log_file=None,
                metrics=metrics
            )
            
            logger.info(f"PETSc solver execution completed in {duration:.1f}s, success={success}")
            return run_result
            
        except subprocess.TimeoutExpired:
            return RunResult(
                success=False,
                exit_code=-1,
                stdout="",
                stderr="PETSc solver execution timed out",
                duration=time.time() - start_time,
                log_file=None,
                metrics={}
            )
        except Exception as e:
            return RunResult(
                success=False,
                exit_code=-1,
                stdout="",
                stderr=f"PETSc solver execution failed: {e}",
                duration=time.time() - start_time,
                log_file=None,
                metrics={}
            )
    
    def check_environment(self) -> EnvironmentResult:
        """Validate runtime environment."""
        missing_components = []
        env_info = {}
        
        try:
            # Check for mpiexec/mpirun
            mpi_cmd = None
            for cmd in ['mpiexec', 'mpirun']:
                result = subprocess.run([cmd, '--version'], capture_output=True, text=True, timeout=10)
                if result.returncode == 0:
                    env_info['mpi_command'] = cmd
                    env_info['mpi_version'] = result.stdout.split('\n')[0]
                    mpi_cmd = cmd
                    break
            
            if not mpi_cmd:
                missing_components.append('mpiexec/mpirun')
            
            # Check PETSc environment
            petsc_dir = os.environ.get('PETSC_DIR')
            if petsc_dir:
                env_info['petsc_dir'] = petsc_dir
                
                # Check if PETSc lib directory exists
                petsc_arch = os.environ.get('PETSC_ARCH', '')
                lib_dir = Path(petsc_dir) / petsc_arch / "lib"
                if not lib_dir.exists():
                    lib_dir = Path(petsc_dir) / "lib"  # Try without PETSC_ARCH
                
                if lib_dir.exists():
                    env_info['petsc_lib'] = str(lib_dir)
                else:
                    missing_components.append('petsc_lib')
            else:
                missing_components.append('PETSC_DIR')
            
            # Check for PETSc arch
            petsc_arch = os.environ.get('PETSC_ARCH')
            if petsc_arch:
                env_info['petsc_arch'] = petsc_arch
            
            valid = len(missing_components) == 0
            
            return EnvironmentResult(
                valid=valid,
                missing_components=missing_components,
                environment_info=env_info,
                error_message=None if valid else f"Missing components: {', '.join(missing_components)}"
            )
            
        except Exception as e:
            return EnvironmentResult(
                valid=False,
                missing_components=['unknown'],
                environment_info={},
                error_message=f"Environment check failed: {e}"
            )
    
    def prepare_environment(self, run_config: RunConfig) -> bool:
        """Prepare environment for execution."""
        try:
            # Check if binary exists
            if not run_config.binary_path.exists():
                logger.error(f"PETSc binary not found: {run_config.binary_path}")
                return False
            
            # Make sure binary is executable
            run_config.binary_path.chmod(0o755)
            
            # Validate PETSc environment
            env_result = self.check_environment()
            if not env_result.valid:
                logger.error(f"Environment validation failed: {env_result.error_message}")
                return False
            
            return True
            
        except Exception as e:
            logger.error(f"Failed to prepare PETSc environment: {e}")
            return False
    
    def cleanup_environment(self, run_config: RunConfig) -> bool:
        """Cleanup after execution."""
        # No specific cleanup needed for PETSc
        return True
    
    def _build_mpiexec_command(self, run_config: RunConfig) -> List[str]:
        """Build the mpiexec command for executing PETSc solver."""
        # Choose mpiexec or mpirun
        mpi_cmd = 'mpiexec'
        try:
            subprocess.run([mpi_cmd, '--version'], capture_output=True, timeout=5)
        except:
            mpi_cmd = 'mpirun'
        
        cmd = [
            mpi_cmd,
            "-n", str(run_config.mpi_processes),
            str(run_config.binary_path)
        ]
        
        # Add PETSc-specific options
        petsc_options = []
        
        # Set solver type
        if run_config.solver_type == 'gcr':
            petsc_options.extend(['-ksp_type', 'gcr'])
        elif run_config.solver_type == 'gmres':
            petsc_options.extend(['-ksp_type', 'gmres'])
        elif run_config.solver_type == 'cg':
            petsc_options.extend(['-ksp_type', 'cg'])
        elif run_config.solver_type == 'bicg':
            petsc_options.extend(['-ksp_type', 'bicg'])
        
        # Set convergence tolerance (default PETSc tolerance)
        petsc_options.extend(['-ksp_rtol', '1e-6'])
        
        # Set maximum iterations
        petsc_options.extend(['-ksp_max_it', '1000'])
        
        # Enable monitoring if in debug mode
        if run_config.debug_mode:
            petsc_options.extend(['-ksp_monitor', '-ksp_max_it', '10'])
        
        # Add PETSc options to command
        cmd.extend(petsc_options)
        
        # Add extra arguments
        cmd.extend(run_config.extra_args)
        
        return cmd
    
    def _extract_petsc_metrics(self, stdout: str, stderr: str) -> Dict[str, Any]:
        """Extract PETSc-specific performance and convergence metrics."""
        metrics = {}
        
        try:
            # Combine stdout and stderr for analysis
            output = stdout + "\n" + stderr
            lines = output.split('\n')
            
            # Extract PETSc KSP solver information
            for line in lines:
                line_lower = line.strip().lower()
                
                # KSP iteration count
                if 'linear solve converged' in line_lower:
                    # Look for "Linear solve converged due to CONVERGED_RTOL iterations 42"
                    parts = line.split()
                    for i, part in enumerate(parts):
                        if part == 'iterations' and i + 1 < len(parts):
                            try:
                                metrics['ksp_iterations'] = int(parts[i + 1])
                            except:
                                pass
                
                # KSP residual norm
                if 'ksp residual norm' in line_lower:
                    # Extract the numerical value
                    parts = line.split()
                    for part in parts:
                        try:
                            if 'e-' in part or 'e+' in part or ('.' in part and part.replace('.', '').replace('-', '').replace('+', '').isdigit()):
                                metrics['final_residual_norm'] = float(part)
                                break
                        except:
                            continue
                
                # Initial residual
                if line_lower.startswith('0 ksp residual norm'):
                    parts = line.split()
                    if len(parts) >= 4:
                        try:
                            metrics['initial_residual_norm'] = float(parts[3])
                        except:
                            pass
                
                # Convergence reason
                if 'converged_rtol' in line_lower:
                    metrics['convergence_reason'] = 'CONVERGED_RTOL'
                    metrics['converged'] = True
                elif 'converged_atol' in line_lower:
                    metrics['convergence_reason'] = 'CONVERGED_ATOL'
                    metrics['converged'] = True
                elif 'diverged' in line_lower:
                    if 'max_it' in line_lower:
                        metrics['convergence_reason'] = 'DIVERGED_MAX_IT'
                    else:
                        metrics['convergence_reason'] = 'DIVERGED'
                    metrics['converged'] = False
                
                # Timing information (if available)
                if 'time' in line_lower and ('sec' in line_lower or 'second' in line_lower):
                    parts = line.split()
                    for part in parts:
                        try:
                            if '.' in part and part.replace('.', '').isdigit():
                                metrics['solver_time'] = float(part)
                                break
                        except:
                            continue
                
                # Matrix and vector info
                if 'matrix type:' in line_lower:
                    parts = line.split(':')
                    if len(parts) > 1:
                        metrics['matrix_type'] = parts[1].strip()
                
                if 'preconditioner type:' in line_lower:
                    parts = line.split(':')
                    if len(parts) > 1:
                        metrics['preconditioner_type'] = parts[1].strip()
        
        except Exception as e:
            logger.debug(f"Failed to extract PETSc metrics: {e}")
        
        return metrics